/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lyh.api.bytecode;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.function.Function;
import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;
import javassist.NotFoundException;
import javassist.bytecode.AnnotationsAttribute;
import javassist.bytecode.ClassFile;
import javassist.bytecode.ConstPool;
import javassist.bytecode.ParameterAnnotationsAttribute;
import javassist.bytecode.annotation.BooleanMemberValue;
import javassist.bytecode.annotation.IntegerMemberValue;
import javassist.bytecode.annotation.LongMemberValue;
import javassist.bytecode.annotation.MemberValue;
import javassist.bytecode.annotation.StringMemberValue;

/**
 * @author shoushen.luan
 * @date 2022/2/17 9:42 上午
 */
public class EnhancerMapper {
    private static final ClassPool pool     = ClassPool.getDefault();
    static EnhancerMapper          instance = null;

    public static EnhancerMapper getInstance() {
        if (EnhancerMapper.instance == null) {
            synchronized (EnhancerMapper.class) {
                if (instance == null) {
                    instance = new EnhancerMapper();
                    pool.insertClassPath(new LoaderClassPath(Thread.currentThread().getContextClassLoader()));
                }
            }
        }
        return EnhancerMapper.instance;
    }

    /**
     * 增加Mapper
     *
     * @param originClazz
     * @return
     */
    public Class<?> enhanceMappper(Class<?> originClazz) {
        try {
            String enhanceClassName = originClazz.getName() + "$JAVASSIST";
            Method[] originMethods = originClazz.getDeclaredMethods();
            CtClass originCtClazz = pool.get(originClazz.getName());
            String methodName[] = new String[originCtClazz.getDeclaredMethods().length];
            final CtClass enhanceCtClass = pool.makeInterface(enhanceClassName);
            copyClazzAnnotations(originCtClazz, enhanceCtClass);
            for (int m = 0; m < originCtClazz.getDeclaredMethods().length; m++) {
                CtMethod ctMethod = originCtClazz.getDeclaredMethods()[m];
                methodName[m] = ctMethod.getName();
                enhanceMethod(enhanceCtClass, originMethods, ctMethod);
            }
            Class<?> loadThisClass = enhanceCtClass.toClass();
            // 输出Class类到本地文件(可选)
            enhanceCtClass.writeFile(".");
            return loadThisClass;
        } catch (Throwable t) {
            t.printStackTrace();
            throw new RuntimeException(String.format("enhance class:`%s` fail", originClazz));
        }
    }

    /**
     * 拷贝注解
     *
     * @param originCtClazz
     * @param enhanceCtClass
     * @throws ClassNotFoundException
     */
    private void copyClazzAnnotations(CtClass originCtClazz, CtClass enhanceCtClass) throws ClassNotFoundException {
        ConstPool cp = enhanceCtClass.getClassFile().getConstPool();
        javassist.bytecode.annotation.Annotation[] annotations = new javassist.bytecode.annotation.Annotation[originCtClazz
            .getAnnotations().length + 1];
        int index = 0;
        AnnotationsAttribute clazzAnnotationsAttribute = new AnnotationsAttribute(cp, AnnotationsAttribute.visibleTag);
        for (Object annotation : originCtClazz.getAnnotations()) {
            annotations[index] = MapperAnnotationEnhancer.makeAnnotation((java.lang.annotation.Annotation) annotation,
                cp);
            index++;
        }
        // 自定义注解
        annotations[index] = createAnnotation(cp);
        enhanceCtClass.getClassFile().addAttribute(clazzAnnotationsAttribute);
        clazzAnnotationsAttribute.setAnnotations(annotations);
    }

    /**
     * 创建自定义注解
     *
     * @return
     */
    private javassist.bytecode.annotation.Annotation createAnnotation(ConstPool cp) {
        javassist.bytecode.annotation.Annotation clazzAnnotation = new javassist.bytecode.annotation.Annotation(
            MyAnnotation2.class.getName(), cp);
        clazzAnnotation.addMemberValue("str", new StringMemberValue("参数", cp));
        clazzAnnotation.addMemberValue("num", new IntegerMemberValue(cp, 123456));
        clazzAnnotation.addMemberValue("bool", new BooleanMemberValue(false, cp));
        return clazzAnnotation;
    }

    /**
     * 增强方法
     *
     * @param enhanceClass
     * @param originMethods
     * @param ctMethod
     * @throws NotFoundException
     * @throws CannotCompileException
     */
    private void enhanceMethod(CtClass enhanceClass, Method[] originMethods, CtMethod ctMethod)
      throws NotFoundException, CannotCompileException {
    String tableSuffix = "$";
    CtMethod newMethod =
        new CtMethod(
            ctMethod.getReturnType(),
            ctMethod.getName() + tableSuffix,
            ctMethod.getParameterTypes(),
            enhanceClass);
    Method method =
        getOriginMethod(
            originMethods, (m) -> newMethod.getName().equals(m.getName() + tableSuffix));
    if (method.getParameterCount() > 0 && method.getParameterAnnotations()[0].length > 0) {
      ClassFile ccFile = enhanceClass.getClassFile();
      ConstPool constPool = ccFile.getConstPool();
      // 拷贝注解信息和注解内容，以支持mybatis mapper类的动态绑定
      ParameterAnnotationsAttribute paa =
          MapperAnnotationEnhancer.duplicateParamAnnoAttribute(constPool, method);
      newMethod.getMethodInfo().addAttribute(paa);
    }
    enhanceClass.addMethod(newMethod);
  }

    private Method getOriginMethod(Method[] originMethods, Function<Method, Boolean> lookupFunc) {
        for (Method method : originMethods) {
            if (lookupFunc.apply(method)) {
                return method;
            }
        }
        throw new RuntimeException("enhanceClass find method error");
    }

    static class MapperAnnotationEnhancer {

        /**
         * 复制方法参数注解
         *
         * @param cp
         * @param method
         * @return
         */
        public static ParameterAnnotationsAttribute duplicateParamAnnoAttribute(ConstPool cp, Method method) {
            ParameterAnnotationsAttribute oldAns = new ParameterAnnotationsAttribute(cp,
                ParameterAnnotationsAttribute.visibleTag);
            int paramAnnoLength = method.getParameterAnnotations().length;
            javassist.bytecode.annotation.Annotation[][] anAr = new javassist.bytecode.annotation.Annotation[paramAnnoLength][];
            for (int i = 0; i < anAr.length; ++i) {
                anAr[i] = new javassist.bytecode.annotation.Annotation[method.getParameterAnnotations()[i].length];
                for (int j = 0; j < anAr[i].length; ++j) {
                    anAr[i][j] = makeAnnotation(method.getParameterAnnotations()[i][j], cp);
                }
            }
            oldAns.setAnnotations(anAr);
            return oldAns;
        }

        /**
         * 创建javassist注解
         *
         * @param annotation
         * @param cp
         * @return
         */
        static javassist.bytecode.annotation.Annotation makeAnnotation(java.lang.annotation.Annotation annotation,
                                                                       ConstPool cp) {
            try {
                Class<? extends Annotation> annoType = annotation.annotationType();
                javassist.bytecode.annotation.Annotation newAnnotation = new javassist.bytecode.annotation.Annotation(
                    annoType.getName(), cp);
                for (Method method : annoType.getDeclaredMethods()) {
                    Object val = method.invoke(annotation);
                    MemberValue memberValue = createMemberValue(method.getReturnType(), val, cp);
                    newAnnotation.addMemberValue(method.getName(), memberValue);
                }
                return newAnnotation;
            } catch (Exception e) {
                throw new RuntimeException("createJavassistAnnotation error", e);
            }
        }

        /**
         * 创建注解属性
         *
         * @param type
         * @param val
         * @param cp
         * @return
         */
        static MemberValue createMemberValue(Class<?> type, Object val, ConstPool cp) {
            if (type == String.class) {
                return new StringMemberValue((String) val, cp);
            } else if (type == Integer.class) {
                return new IntegerMemberValue((Integer) val, cp);
            } else if (type == Long.class) {
                return new LongMemberValue((Long) val, cp);
            } else if (type == Boolean.class) {
                return new BooleanMemberValue((Boolean) val, cp);
            } else { // ...
                throw new RuntimeException(String.format("NoSupport param type:`%s`,value:`%s`", type, val));
            }
        }
    }
}
